clear; clc
%%
tau = 1;%s
gain = -3.0e3; %degree/Hz
order = 8; bw = 20; %Hz


g=nmr_osc(tau, gain); lp=lpFilter(order, bw);
plant = g*lp;
% ax1 = subplot(2, 1, 1); hold off;
% ax2 = subplot(2, 1, 2); hold off;
% 
% xList=0.01:0.03:0.1;
% for x=xList
%     rt_x= referenc_tracking(gain, tau, g, lp, x); 
%     nr_x = noise_rejection(gain, tau, g, lp, x); 
%     stepplot(ax1, rt_x); hold on;
%     stepplot(ax2, nr_x); hold on;
% end
% legend(ax1, arrayfun(@(x) sprintf('x=%f', x), xList, 'UniformOutput', false));
% legend(ax2, arrayfun(@(x) sprintf('x=%f', x), xList, 'UniformOutput', false));

%%
seed=1; rate=100; totT=60.0; ampP=0; ampW=0.005; 
x=0.05;

pxxN = []; pxxE = [];
for seed = 1:100
    [modelO, modeC, modelE, kp, ki, noise, yo, yc, ye, t, pxxN1, pxxE1, f] = sim(seed, rate, totT, ampP, ampW, x, tau, gain, order, bw, g, lp);
    pxxN=[pxxN pxxN1]; pxxE=[pxxE pxxE1]; 
end
%%
figure;
subplot(3, 1, 1)
plot(t, yo); ylabel('phase (deg)')
subplot(3, 1, 2)
plot(t, yc); ylabel('phase (deg)'); legend(sprintf('Kp=%3.2fm\nKi=%3.2fm', kp*1000, ki*1000));
subplot(3, 1, 3)
plot(t, noise*1e3, t, ye*1e3); ylabel('freqeuncy (mHz)'); legend('noise', 'pid output')

figure
loglog(f, sqrt(mean(pxxN,2)), f, mean(sqrt(pxxE), 2) ); grid on; ylim([1e-4 1e-2])
%%

function [modelO, modelC, modelE, kp, ki, noise,  yo, yc, ye, t, pxxN, pxxE, f] = sim(seed, rate, totT, ampP, ampW, x, tau, gain, order, bw, g, lp)
    rng(seed)
    nt=totT*rate; tList = 0:(nt-1); 
    
    cnP = dsp.ColoredNoise('Color', 'pink', 'SamplesPerFrame', nt); noiseP = cnP(); 
    cnW = dsp.ColoredNoise('Color', 'white', 'SamplesPerFrame', nt); noiseW = cnW();
    noise=ampP*noiseP+ampW*noiseW; 

    modelO = nmr_osc(tau, gain)*lpFilter(order, bw);
    [modelC, kp, ki] = noise_rejection(gain, tau, g, lp, x);
    modelE = control_effort(gain, tau, g, lp, x);

    yo = lsim(modelO, noise, tList/rate);
    yc = lsim(modelC, noise, tList/rate);
    [ye, t] = lsim(modelE, noise, tList/rate);

%     figure;
%     subplot(3, 1, 1)
%     plot(t, yo); ylabel('phase (deg)')
%     subplot(3, 1, 2)
%     plot(t, yc); ylabel('phase (deg)'); legend(sprintf('Kp=%3.2fm\nKi=%3.2fm', kp*1000, ki*1000));
%     subplot(3, 1, 3)
%     plot(t, noise*1e3, t, ye*1e3); ylabel('freqeuncy (mHz)'); legend('noise', 'pid output')


%     figure;
    [pxxN, ~] = periodogram(noise, rectwin(length(ye)),length(ye),rate);
    [pxxE, f] = periodogram(ye, rectwin(length(ye)),length(ye),rate);
%     loglog(f, sqrt(pxxN), f, sqrt(pxxE)); grid on; ylim([1e-4 1e-2])
end

function plant = nmr_osc(tau, gain)
    s = tf('s');
    plant = gain / (1 + tau*s);
end

function model = lpFilter(order, bw)
    s=tf('s');
    model = (1/(1+s/(2*pi*bw)))^order;
end

function c = controller(gain, tau, x)
    x1 = min([0.25 x]);
    kp = 1.0./gain./x; 
    tauI = 4.0*x1*tau; 
    ki = kp/tauI;
    c = pid(kp, ki);
end

function [rt, kp, ki] = referenc_tracking(gain, tau, plant, filter, x)
    g=plant*filter;
    c=controller(gain, tau, x);
    rt = feedback(g*c, 1);
    kp=c.Kp; ki=c.Ki;
end

function [nr, kp, ki] = noise_rejection(gain, tau, plant, filter, x)
    g=plant*filter;
    c=controller(gain, tau, x);
    nr = feedback(g, c);
    kp=c.Kp; ki=c.Ki;
end

function [ce, kp, ki] = control_effort(gain, tau, plant, filter, x)
    c=controller(gain, tau, x);
    [nr, kp, ki] = noise_rejection(gain, tau, plant, filter, x);
    ce = nr*c;
end